import cv2
import numpy as np
import torchvision

def predict(left_img, right_img):
    left_img = cv2.cvtColor(np.array(torchvision.transforms.ToPILImage()(left_img).convert('RGB')), cv2.COLOR_RGB2BGR)
    right_img = cv2.cvtColor(np.array(torchvision.transforms.ToPILImage()(right_img).convert('RGB')), cv2.COLOR_RGB2BGR)

    stereo = cv2.StereoSGBM_create()

    # 设置参数
    # kitti
    # stereo.setMinDisparity(1) # 最小视差
    # stereo.setNumDisparities(192) # 最大视差
    # stereo.setBlockSize(9) # 块(邻域)大小
    # stereo.setP1(16 * 3 * stereo.getBlockSize() ** 2) # 惩罚参数，用于平滑视差
    # stereo.setP2(32 * 3 * stereo.getBlockSize() ** 2)
    # stereo.setUniquenessRatio(15) # 唯一性比率
    # stereo.setSpeckleWindowSize(1) # 去噪窗口尺寸
    # stereo.setSpeckleRange(10) # 去噪范围

    # usvinland
    stereo.setMinDisparity(1) # 最小视差
    stereo.setNumDisparities(64) # 最大视差
    stereo.setBlockSize(15) # 块(邻域)大小
    stereo.setP1(16 * 3 * stereo.getBlockSize() ** 2) # 惩罚参数，用于平滑视差
    stereo.setP2(32 * 3 * stereo.getBlockSize() ** 2)

    # 估计视差
    disp = np.uint8(stereo.compute(left_img, right_img))

    return disp
